import time
import math
import serial
import struct
import threading
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from threading import Lock

class IMUDriver(Node):
    def __init__(self):
        super().__init__('imu_driver')
        
        # 声明参数
        self.declare_parameter('port', '/dev/ttyUSB0')
        self.declare_parameter('baudrate', 115200)
        
        # 获取参数
        port = self.get_parameter('port').value
        baudrate = self.get_parameter('baudrate').value
        
        # IMU数据存储
        self.angular_velocity = [0.0, 0.0, 0.0]
        self.acceleration = [0.0, 0.0, 0.0]
        self.angle_degree = [0.0, 0.0, 0.0]
        self.quaternion = [1.0, 0.0, 0.0, 0.0]
        self.data_lock = Lock()
        
        # 创建IMU消息发布器
        self.imu_pub = self.create_publisher(Imu, 'imu/data', 10)
        
        # 初始化串口
        self.ser = self.init_serial(port, baudrate)
        
        # 启动串口读取线程
        self.read_thread = threading.Thread(target=self.serial_read_loop)
        self.read_thread.daemon = True
        self.read_thread.start()
        
        # 定时发布IMU数据
        self.timer = self.create_timer(0.01, self.publish_imu_data)
        
        self.get_logger().info(f"IMU driver started on {port} with baudrate {baudrate}")
    
    def init_serial(self, port, baudrate):
        """初始化串口连接"""
        try:
            ser = serial.Serial(port=port, baudrate=baudrate, timeout=0.5)
            if ser.is_open:
                self.get_logger().info("Serial port opened successfully")
                # 清空缓冲区
                ser.reset_input_buffer()
                ser.reset_output_buffer()
                return ser
            else:
                self.get_logger().error("Failed to open serial port")
                return None
        except Exception as e:
            self.get_logger().error(f"Serial port error: {e}")
            return None
    
    def parse_short(self, low_byte, high_byte):
        """解析有符号short类型数据"""
        # 按照协议：DATA = (short)((short)DATAH << 8 | DATAL)
        # 在Python中，使用struct来解析
        value_bytes = bytes([low_byte, high_byte])
        value = struct.unpack('<h', value_bytes)[0]  # 小端序
        return value
    
    def calculate_checksum(self, data_list):
        """计算校验和"""
        return sum(data_list) & 0xFF
    
    def parse_data_packet(self, packet):
        """解析完整的数据包"""
        if len(packet) != 11:  # 0x55 + TYPE + 8数据字节 + 1校验和
            return False, None
        
        start_byte = packet[0]
        data_type = packet[1]
        data_bytes = packet[2:10]  # 8个数据字节
        received_checksum = packet[10]
        
        if start_byte != 0x55:
            return False, None
        
        # 计算校验和（包括0x55和TYPE）
        checksum_data = [start_byte, data_type] + list(data_bytes)
        calculated_checksum = self.calculate_checksum(checksum_data)
        
        if calculated_checksum != received_checksum:
            return False, data_type
        
        # 解析数据（4个有符号short值）
        values = []
        for i in range(0, 8, 2):
            if i + 1 < len(data_bytes):
                value = self.parse_short(data_bytes[i], data_bytes[i+1])
                values.append(value)
        
        return True, (data_type, values)
    
    def process_imu_data(self, data_type, values):
        """处理解析后的IMU数据"""
        if len(values) < 3:
            return
        
        try:
            with self.data_lock:
                if data_type == 0x51:  # 加速度数据
                    # 官方转换公式：/ 32768.0 * 16 * 9.8
                    self.acceleration = [x / 32768.0 * 16 * 9.8 for x in values[:3]]
                    self.get_logger().debug(f"Acceleration: {self.acceleration}", throttle_duration_sec=2)
                    
                elif data_type == 0x52:  # 角速度数据
                    # 官方转换公式：/ 32768.0 * 2000 * π / 180
                    self.angular_velocity = [x / 32768.0 * 2000 * math.pi / 180 for x in values[:3]]
                    self.get_logger().debug(f"Angular velocity: {self.angular_velocity}", throttle_duration_sec=2)
                    
                elif data_type == 0x53:  # 角度数据
                    # 官方转换公式：/ 32768.0 * 180
                    self.angle_degree = [x / 32768.0 * 180 for x in values[:3]]
                    self.get_logger().debug(f"Angle: {self.angle_degree}", throttle_duration_sec=2)
                    
                elif data_type == 0x54:  # 磁场数据
                    magnetometer = [x for x in values[:3]]
                    self.get_logger().debug(f"Magnetometer: {magnetometer}", throttle_duration_sec=2)
                    
                elif data_type == 0x59:  # 四元数数据
                    # 四元数归一化：/ 32768.0
                    self.quaternion = [x / 32768.0 for x in values[:4]]
                    self.get_logger().debug(f"Quaternion: {self.quaternion}", throttle_duration_sec=2)
                    
        except Exception as e:
            self.get_logger().error(f"Error processing IMU data: {e}")
    
    def serial_read_loop(self):
        """串口读取线程"""
        if not self.ser:
            return
        
        buffer = bytearray()
        packet_length = 11  # 完整数据包长度
        
        self.get_logger().info("Starting serial read loop")
        
        while rclpy.ok():
            try:
                # 读取可用数据
                if self.ser.in_waiting > 0:
                    data = self.ser.read(self.ser.in_waiting)
                    buffer.extend(data)
                    
                    # 处理缓冲区中的数据
                    while len(buffer) >= packet_length:
                        # 查找起始字节0x55
                        start_index = -1
                        for i in range(len(buffer) - packet_length + 1):
                            if buffer[i] == 0x55:
                                start_index = i
                                break
                        
                        if start_index == -1:
                            # 没有找到起始字节，清空缓冲区
                            buffer.clear()
                            break
                        
                        # 提取完整数据包
                        if start_index > 0:
                            # 丢弃起始字节之前的数据
                            buffer = buffer[start_index:]
                        
                        if len(buffer) >= packet_length:
                            packet = buffer[:packet_length]
                            buffer = buffer[packet_length:]
                            
                            # 解析数据包
                            success, result = self.parse_data_packet(packet)
                            if success:
                                data_type, values = result
                                self.process_imu_data(data_type, values)
                            else:
                                if result is not None:  # 校验和错误
                                    data_type = result
                                    self.get_logger().warn(
                                        f"Checksum error for type 0x{data_type:02x}, "
                                        f"packet: {['0x{:02x}'.format(b) for b in packet]}",
                                        throttle_duration_sec=1
                                    )
                else:
                    time.sleep(0.001)  # 短暂休眠，避免CPU占用过高
                    
            except Exception as e:
                self.get_logger().error(f"Serial read error: {e}")
                time.sleep(0.1)
                # 清空缓冲区
                buffer.clear()
    
    def euler_to_quaternion(self, roll, pitch, yaw):
        """将欧拉角转换为四元数"""
        cy = math.cos(yaw * 0.5)
        sy = math.sin(yaw * 0.5)
        cp = math.cos(pitch * 0.5)
        sp = math.sin(pitch * 0.5)
        cr = math.cos(roll * 0.5)
        sr = math.sin(roll * 0.5)
        
        qw = cr * cp * cy + sr * sp * sy
        qx = sr * cp * cy - cr * sp * sy
        qy = cr * sp * cy + sr * cp * sy
        qz = cr * cp * sy - sr * sp * cy
        
        return [qx, qy, qz, qw]
    
    def publish_imu_data(self):
        """发布IMU数据"""
        with self.data_lock:
            imu_msg = Imu()
            imu_msg.header.stamp = self.get_clock().now().to_msg()
            imu_msg.header.frame_id = 'imu_link'
            
            # 设置线性加速度 (m/s²)
            imu_msg.linear_acceleration.x = self.acceleration[0]
            imu_msg.linear_acceleration.y = self.acceleration[1]
            imu_msg.linear_acceleration.z = self.acceleration[2]
            
            # 设置角速度 (rad/s)
            imu_msg.angular_velocity.x = self.angular_velocity[0]
            imu_msg.angular_velocity.y = self.angular_velocity[1]
            imu_msg.angular_velocity.z = self.angular_velocity[2]
            
            # 设置方向（优先使用四元数，如果没有则使用角度转换）
            if max(abs(x) for x in self.quaternion) > 0.1:  # 检查四元数是否有效
                imu_msg.orientation.x = self.quaternion[1]
                imu_msg.orientation.y = self.quaternion[2]
                imu_msg.orientation.z = self.quaternion[3]
                imu_msg.orientation.w = self.quaternion[0]
            else:
                # 使用角度转换为四元数
                roll = math.radians(self.angle_degree[0])
                pitch = math.radians(self.angle_degree[1])
                yaw = math.radians(self.angle_degree[2])
                quaternion = self.euler_to_quaternion(roll, pitch, yaw)
                imu_msg.orientation.x = quaternion[0]
                imu_msg.orientation.y = quaternion[1]
                imu_msg.orientation.z = quaternion[2]
                imu_msg.orientation.w = quaternion[3]
            
            # 设置协方差矩阵
            imu_msg.orientation_covariance[0] = -1
            imu_msg.angular_velocity_covariance[0] = -1
            imu_msg.linear_acceleration_covariance[0] = -1
            
            self.imu_pub.publish(imu_msg)
            
            # # 限制日志输出频率
            # self.get_logger().info("IMU data published", throttle_duration_sec=1)

def main(args=None):
    rclpy.init(args=args)
    imu_driver = IMUDriver()
    
    try:
        rclpy.spin(imu_driver)
    except KeyboardInterrupt:
        pass
    except Exception as e:
        imu_driver.get_logger().error(f"Error: {e}")
    finally:
        if imu_driver.ser and imu_driver.ser.is_open:
            imu_driver.ser.close()
        imu_driver.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()